/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file fused_add_topk_div_case.cpp
 * \brief FusedAddTopkDiv 测试用例.
 */
#include "fused_add_topk_div_case.h"
#include <utility>
#include <tikicpulib.h>
#include <graph/utils/type_utils.h>
#include <register/op_impl_registry.h>
#include "tests/utils/log.h"
#include "tests/utils/platform.h"
#include "tiling/tiling_base.h"

/**
 * 以下函数声明需要保持与 CMakeList.txt 中调用 OpsTest_Level2_AddOp 函数时 KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT
 * 参数所控制的 Kernel 入口一致.
 */

#define FUSED_ADD_TOPK_DIV_KERNEL_PARAM                                                                               \
    (GM_ADDR x, GM_ADDR addNum, GM_ADDR mappingNum, GM_ADDR mappingTable, GM_ADDR y, GM_ADDR indices,                 \
     GM_ADDR workspace, GM_ADDR tiling)

using FusedAddTopkDivKernelFunc = void(*) FUSED_ADD_TOPK_DIV_KERNEL_PARAM;

extern "C" __global__ __aicore__ void fused_add_topk_div FUSED_ADD_TOPK_DIV_KERNEL_PARAM;

using namespace ops::adv::tests::FusedAddTopkDiv;
using TensorIntf = ops::adv::tests::utils::TensorIntf;
using Case = ops::adv::tests::utils::Case;
using Platform = ops::adv::tests::utils::Platform;

bool RunFusedAddTopkDiv(void *func, uint64_t tilingKey, int64_t blockDim, std::vector<TensorIntf *> &inputs,
                        std::vector<TensorIntf *> &outputs, uint8_t *workspace, uint8_t *tilingData)
{
    // Kernel 运行
    auto kernelFunc = (FusedAddTopkDivKernelFunc)func;
    ICPU_SET_TILING_KEY(tilingKey);
    ICPU_RUN_KF(kernelFunc, blockDim, inputs[X_INDEX]->GetDevData(), inputs[ADD_NUM_INDEX]->GetDevData(),
                inputs[MAPPING_NUM_INDEX]->GetDevData(), inputs[MAPPING_TABLE_INDEX]->GetDevData(),
                outputs[Y_INDEX]->GetDevData(), outputs[INDICES_INDEX]->GetDevData(), workspace, tilingData);
    return true;
}

extern "C" ge::graphStatus TilingForFusedAddTopkDivStub(gert::TilingContext *context)
{
    auto *fusedAddTopkDivCase = static_cast<FusedAddTopkDivCase *>(Case::GetCurrentCase());
    if (fusedAddTopkDivCase != nullptr) {
        FusedAddTopkDivCase::DoTilingParam p;
        p.ctx = context;
        p.ret = ge::GRAPH_SUCCESS;
        if (!fusedAddTopkDivCase->DoOpTiling(p)) {
            return p.ret;
        }
        return fusedAddTopkDivCase->fusedAddTopkDivTilingFunc(context);
    }
    return ge::GRAPH_FAILED;
}

bool FusedAddTopkDivCase::InitParam()
{
    x = Tensor("x", {mParam.a, mParam.b}, "2", mParam.xDtype, ge::FORMAT_ND);
    addNum = Tensor("addNum", {mParam.b}, "1", mParam.xDtype, ge::FORMAT_ND);
    mappingNumOptional = Tensor("mappingNumOptional", {mParam.b}, "1", ge::DataType::DT_INT32, ge::FORMAT_ND,
                                Tensor::TensorType::OPTIONAL_INPUT);
    mappingTableOptional = Tensor("mappingTableOptional", {mParam.b, mParam.c}, "2", ge::DataType::DT_INT32,
                                  ge::FORMAT_ND, Tensor::TensorType::OPTIONAL_INPUT);

    yOut = Tensor("yOut", {mParam.a, mParam.k}, "2", ge::DataType::DT_FLOAT, ge::FORMAT_ND);
    indicesOut = Tensor("indicesOut", {mParam.a, mParam.k}, "2", ge::DataType::DT_INT32, ge::FORMAT_ND);
    return true;
}

bool FusedAddTopkDivCase::InitOpInfo()
{
    bool rst = mCtx.SetOpName("FusedAddTopkDiv");
    rst = rst && mCtx.SetDeterministic(false);
    rst = rst && mCtx.SetInputs({&x, &addNum, &mappingNumOptional, &mappingTableOptional});
    rst = rst && mCtx.SetOutputs({&yOut, &indicesOut});
    rst = rst && mCtx.SetAttrs({{"group_num", mParam.groupNum},
                                {"group_topk", mParam.groupTopk},
                                {"n", mParam.n},
                                {"k", mParam.k},
                                {"activate_type", mParam.activateType},
                                {"is_norm", mParam.isNorm},
                                {"enable_expert_mapping", mParam.enableExpertMapping},
                                {"scale", mParam.scale}});
    rst = rst && mCtx.SetKernelRunCbf(RunFusedAddTopkDiv);
    rst = rst && mCtx.SetKernelMainFunc((void *)fused_add_topk_div);
    rst = rst && mOpInfo.SetContext(&mCtx);

    auto *platform = Platform::GetGlobalPlatform();
    if (platform == nullptr) {
        LOG_ERR("Global Platform is null");
        return false;
    }

    fusedAddTopkDivTilingFunc =
        (gert::OpImplRegisterV2::TilingKernelFunc)platform->LoadOpTilingSoSym("TilingForFusedAddTopkDiv");
    if (fusedAddTopkDivTilingFunc == nullptr) {
        LOG_ERR("Can't get origin tiling func, fusedAddTopkDiv(%p)", fusedAddTopkDivTilingFunc);
        return false;
    }
    IMPL_OP(FusedAddTopkDiv).Tiling(TilingForFusedAddTopkDivStub);
    return rst;
}

bool FusedAddTopkDivCase::InitCurrentCasePtr()
{
    Case::mCurrentCasePtr = this;
    return true;
}

bool FusedAddTopkDivCase::Run()
{
    if (!mEnable) {
        return true;
    }
    if (!mOpInfo.ProcessTiling(mName)) {
        return false;
    }
    if (!mOpInfo.ProcessKernel(mName)) {
        return false;
    }
    return true;
}

FusedAddTopkDivCase::FusedAddTopkDivCase(const char *name, bool enable, const char *dbgInfo, OpInfo incre,
                                           Param param)
    : Case(name, enable, dbgInfo), mOpInfo(std::move(incre)), mParam(std::move(param))
{
    this->mOpInfo.mName = "FusedAddTopkDiv";
}

FusedAddTopkDivCase::FusedAddTopkDivCase()
{
}

bool FusedAddTopkDivCase::DoOpTiling(DoTilingParam &tilingParam)
{
    if (tilingParam.ctx == nullptr) {
        return false;
    }
    return true;
}